package com.yvan.dsf.tracing.client;

import com.yvan.dsf.tracing.api.Span;

import java.util.EmptyStackException;
import java.util.Stack;

public class ContextHolder {

    //一个线程中span栈,用来存放一个线程中多个span情况
    private static ThreadLocal<Stack<Span>> localSpan = new ThreadLocal<Stack<Span>>() {
        @Override
        protected Stack<Span> initialValue() {
            return new Stack<Span>();
        }
    };

    private static ThreadLocal<String> localTraceId = new ThreadLocal<String>();

    private static ThreadLocal<Boolean> localSample = new ThreadLocal<Boolean>() {
        @Override
        protected Boolean initialValue() {
            return true;
        }
    };

    public  static String getTraceId() {
        return localTraceId.get();
    }

    public static void setTraceId(String traceId) {
        localTraceId.set(traceId);
    }

    public static void removeTraceId() {
        localTraceId.remove();
    }

    public static boolean isSample() {
        return localSample.get();
    }

    public static void setLocalSample(boolean isSample) {
        localSample.set(isSample);
    }

    public static void removeSample() {
        localSample.remove();
    }

    public static void setSpan(Span span) {
        ContextHolder.localSpan.get().push(span);
    }

    public static Span getSpan() {
        try {
            return localSpan.get().peek();
        } catch (EmptyStackException e) {
            return null;
        }
    }

    public static Span popSpan() {
        try {
            return localSpan.get().pop();
        } catch (EmptyStackException e) {
            return null;
        }
    }


    public static void removeSpan() {
        localSpan.remove();
    }

    public static void removeAll() {
        if (localSpan.get().size() <= 0) {//span堆栈为空的时候才能清除
            removeSpan();
            removeTraceId();
            removeSample();
        }
    }


}
